from flask import Flask, request, Response
from flask_cors import CORS  # 新增CORS支持
from openai import OpenAI
import json

app = Flask(__name__)
CORS(app)  # 启用跨域支持

# 初始化客户端（与原代码一致）
client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")

@app.route('/chat', methods=['POST'])
def chat_stream():
    # 获取客户端传递的对话历史
    request_data = request.get_json()
    messages = request_data.get('messages', [])

    # 创建流式对话补全
    completion = client.chat.completions.create(
        model="local-model",
        messages=messages,
        temperature=0.7,
        stream=True
    )

    def event_stream():
        new_message = {"role": "assistant", "content": ""}
        # 流式处理每个数据块
        for chunk in completion:
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                new_message["content"] += content
                # 实时返回每个token
                # 确保每个数据块都符合SSE格式
                yield f"data: {json.dumps({'token': content})}\n\n"

    # 保持返回Response对象
    return Response(event_stream(), mimetype='text/event-stream')

if __name__ == '__main__':
    app.run(port=8001)
